from collections import defaultdict
from pathlib import Path
from tqdm import tqdm
import argparse
import numpy as np
import json
from sklearn.model_selection import KFold

from utils import process_record, save_record, generate_stratified_split, load_patient_data, get_distribution


DATA_ROOT = Path('/data/IDLab/DigiHealth/st-petersburg-incart-12-lead-arrhythmia-database-1.0.0/files')
PATIENT_MAPPING = {
    1: ['I01', 'I02'],
    2: ['I03', 'I04', 'I05'],
    3: ['I06', 'I07'],
    4: ['I08'],
    5: ['I09', 'I10', 'I11'],
    6: ['I12', 'I13', 'I14'],
    7: ['I15'],
    8: ['I16', 'I17'],
    9: ['I18', 'I19'],
    10: ['I20', 'I21', 'I22'],
    11: ['I23', 'I24'],
    12: ['I25', 'I26'],
    13: ['I27', 'I28'],
    14: ['I29', 'I30', 'I31', 'I32'],
    15: ['I33', 'I34'],
    16: ['I35', 'I36', 'I37'],
    17: ['I38', 'I39'],
    18: ['I40', 'I41'],
    19: ['I42', 'I43'],
    20: ['I44', 'I45', 'I46'],
    21: ['I47', 'I48'],
    22: ['I49', 'I50'],
    23: ['I51', 'I52', 'I53'],
    24: ['I54', 'I55', 'I56'],
    25: ['I57', 'I58'],
    26: ['I59', 'I60', 'I61'],
    27: ['I62', 'I63', 'I64'],
    28: ['I65', 'I66', 'I67'],
    29: ['I68', 'I69'],
    30: ['I70', 'I71'],
    31: ['I72', 'I73'],
    32: ['I74', 'I75']
}
LABEL_MAPPING = {
    'N': 0, 'j': 0, 'B': 0, 'R': 0, # N
    'A': 1, 'S': 1, 'n': 1, # SVEB
    'V': 2, # VEB
    'F': 3, # F
    'Q': 4 # Q
}

WINDOW_SIZE = 1000 # in # of samples
SAMPLING_RATE = 257
NEW_SAMPLING_RATE = None
TIME_ENCODED = False
OUTPUT_FOLDER = Path(f'/data/IDLab/DigiHealth/processed_data/beat-to-beat/incart-time')
SPLIT_FOLDER = Path(f'/home/timodw/IDLab/Digihealth-Asia/cardiovascular_monitoring/ugent/heartbeat_classification/processed_data/configs')

TEST_RATIO = 1 / 2

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--stratified', action='store_true')
    args = parser.parse_args()

    if args.stratified:
        training_ids, testing_ids = generate_stratified_split(OUTPUT_FOLDER, PATIENT_MAPPING.keys(), testing_fraction=TEST_RATIO)
        evaluation_ids, testing_ids = generate_stratified_split(OUTPUT_FOLDER, testing_ids, testing_fraction=.5)
        training_distribution = get_distribution(OUTPUT_FOLDER, training_ids.tolist())
        evaluation_distribution = get_distribution(OUTPUT_FOLDER, evaluation_ids.tolist())
        testing_distribution = get_distribution(OUTPUT_FOLDER, testing_ids.tolist())
        json_dict = [{
            'training': {
                'distribution': training_distribution,
                'patient_ids': training_ids.tolist()
            },
            'evaluation': {
                'distribution': evaluation_distribution,
                'patient_ids': evaluation_ids.tolist()
            },
            'testing': {
                'distribution': testing_distribution,
                'patient_ids': testing_ids.tolist()
            }
        },]
        json_str = json.dumps(json_dict, sort_keys=True, indent=4)
        print(json_str)
        print(json_str, file=open(SPLIT_FOLDER / f'incart_stratified_standard.json', 'w'))
    else:
        for sr in range(100, 501, 25):
            OUTPUT_FOLDER = Path(f'/data/IDLab/DigiHealth/processed_data/beat-to-beat/incart-time-{sr}')
            print(OUTPUT_FOLDER.name)
            
            for patient_id, file_ids in tqdm(PATIENT_MAPPING.items()):
                for i, f in enumerate(file_ids):
                    result = process_record(DATA_ROOT, f, LABEL_MAPPING,
                                            window_size=WINDOW_SIZE,
                                            sampling_rate=SAMPLING_RATE, lead='II',
                                            resampling=None, time_encoded=True)
                    if result is not None:
                        patient_folder = OUTPUT_FOLDER / f"PATIENT_{patient_id}" / f"RECORD_{i}"
                        save_record(patient_folder, *result)



